from glob import glob

import torch
from tqdm import tqdm
from pathlib import Path
from torch.utils.data import Dataset
import cv2
import numpy as np
import os
from PIL import Image
import random
from datasets.table.augment_table import AugmentBright,Augment
augbright=AugmentBright()
augment=Augment()
def deal_mask(path, width, height):
    img = cv2.imread(path)
    #     img = cv2.resize(img, (width, height))
    kernel = cv2.getStructuringElement(cv2.MORPH_CROSS, (2, 2))
    label = cv2.dilate(img, kernel)
    label[label > 8] = 255
    label[label <= 8] = 0
    return label


def deal_mask1(path, width, height):
    img = cv2.imread(path, 0)

    indexes = np.where(img > 0)
    fx, fy = img.shape[1] / 320, img.shape[0] / 200
    new_indexes = []
    coord_y = np.floor(indexes[0] / fy)
    coord_x = np.floor(indexes[1] / fx)
    new_indexes.append(np.array(coord_y, dtype=np.int64))
    new_indexes.append(np.array(coord_x, dtype=np.int64))
    img[img >= 0] = 0
    img = cv2.resize(img, (width // 4, height // 4))  # , cv2.INTER_AREA)
    img[tuple(new_indexes)] = 255
    return img

def aug_img_function(sourceimg_name,labelimg_name):
    label_im = cv2.imread(labelimg_name)
    source_im = cv2.imread(sourceimg_name)
    image_source_im = Image.fromarray(cv2.cvtColor(source_im, cv2.COLOR_BGR2RGB))
    image_label_im = Image.fromarray(label_im)
    rand_num = random.uniform(0, 1)
    if rand_num < 0.5:
        new_im = augbright.light_exchange(source_im)
        new_label_im = label_im
    else:
        new_im, new_label_im = augment.randomfunction(image_source_im, image_label_im)
        new_im = cv2.cvtColor(np.asarray(new_im), cv2.COLOR_RGB2BGR)
        new_label_im = np.asarray(new_label_im)
    return new_im,new_label_im
class LoadTableImageAndLabels(Dataset):
    def __init__(self, fd):
        self.label_ext = '-label.jpg'
        self.image_ext = '.jpg'
        self.image_ext_png = '.png'
        self.image_ext_jpeg = '.jpeg'
        self.fd = fd
        self.labels = [str(i) for i in Path(fd).glob('*' + self.label_ext)]  # [:100]
        self.images = []
        for i in tqdm(self.labels):
            if os.path.exists(i.replace(self.label_ext, self.image_ext)):
                name = i.replace(self.label_ext, self.image_ext)
            elif os.path.exists(i.replace(self.label_ext, self.image_ext_png)):
                name = i.replace(self.label_ext, self.image_ext_png)
            elif os.path.exists(i.replace(self.label_ext, self.image_ext_jpeg)):
                name = i.replace(self.label_ext, self.image_ext_jpeg)
            else:
                continue
            self.images.append(name)

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, index):
        im,label=aug_img_function(self.images[index],self.labels[index])
        img = im / 255
        img = img[:, :, ::-1].transpose(2, 0, 1)
        img = np.ascontiguousarray(img)

        # label = self.label_images[index]
        ##########
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (7, 7))
        dilateimg = cv2.dilate(label, kernel)
        dilateimg = dilateimg - label
        ###########
        # kernel = np.ones((3, 3), np.uint8)
        # label = cv2.erode(label,kernel,iterations=1)
        # label = cv2.resize(label, (width//4, height//4), cv2.INTER_CUBIC)
        # if index % 10 == 0:
        #    cv2.imwrite('results/gt_'+str(index)+'.jpg', self.label_images[index])
        # print(label.shape)
        # label = self.label_images[index] / 255
        label = label[:, :, 0:1].transpose(2, 0, 1)
        label = np.ascontiguousarray(label) / 255
        dilateimg = dilateimg[:, :, 0:1].transpose(2, 0, 1)
        dilateimg = np.ascontiguousarray(dilateimg) / 255
        return torch.from_numpy(img).float(), torch.from_numpy(label).float(), torch.from_numpy(dilateimg).float()
#数据增强部分在外面做
# class LoadTableImageAndLabels(Dataset):
#     def __init__(self, fd):
#         self.label_ext = '-label.jpg'
#         self.image_ext = '.jpg'
#         self.image_ext_png = '.png'
#         self.image_ext_jpeg = '.jpeg'
#         self.fd = fd
#         self.labels = [str(i) for i in Path(fd).glob('*' + self.label_ext)]  # [:100]
#         width, height = 1024, 640
#         self.label_images = [deal_mask(i, width, height) for i in tqdm(self.labels)]
#         self.images = []
#         for i in tqdm(self.labels):
#             if os.path.exists(i.replace(self.label_ext, self.image_ext)):
#                 name = i.replace(self.label_ext, self.image_ext)
#             elif os.path.exists(i.replace(self.label_ext, self.image_ext_png)):
#                 name = i.replace(self.label_ext, self.image_ext_png)
#             elif os.path.exists(i.replace(self.label_ext, self.image_ext_jpeg)):
#                 name = i.replace(self.label_ext, self.image_ext_jpeg)
#             else:
#                 continue
#             im = cv2.imread(name)
#             self.images.append(im)
#
#     def __len__(self):
#         return len(self.labels)
#
#     def __getitem__(self, index):
#         img = self.images[index] / 255
#         img = img[:, :, ::-1].transpose(2, 0, 1)
#         img = np.ascontiguousarray(img)
#
#         label = self.label_images[index]
#         ##########
#         kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (7, 7))
#         dilateimg = cv2.dilate(label, kernel)
#         dilateimg = dilateimg - label
#         ###########
#         # kernel = np.ones((3, 3), np.uint8)
#         # label = cv2.erode(label,kernel,iterations=1)
#         # label = cv2.resize(label, (width//4, height//4), cv2.INTER_CUBIC)
#         # if index % 10 == 0:
#         #    cv2.imwrite('results/gt_'+str(index)+'.jpg', self.label_images[index])
#         # print(label.shape)
#         # label = self.label_images[index] / 255
#         label = label[:, :, 0:1].transpose(2, 0, 1)
#         label = np.ascontiguousarray(label) / 255
#         dilateimg = dilateimg[:, :, 0:1].transpose(2, 0, 1)
#         dilateimg = np.ascontiguousarray(dilateimg) / 255
#         return torch.from_numpy(img).float(), torch.from_numpy(label).float(), torch.from_numpy(dilateimg).float()
if __name__ == '__main__':
    data_fd = r'F:\datasets\ocr\internal\mask-label'
    dataset = LoadTableImageAndLabels(data_fd)
    dataset[1]
